#!/usr/bin/python
# -*- encoding: utf-8 -*-

"""
@File:      tc_ai_jupyterlab.py
@Time:      2023/07/4 11:12:50
@Author:    liangjian
@Version:   1.0

"""
from common.basetest import LocalTest,BaseTest

from avocado.utils import process
import time,re
#import requests

class Test(BaseTest):
    """
    See tc_ai_pytorch_quici_start.yaml for details

    :avocado: tags=fix,P1,noarch,local,ai_container,jupyterlab_ai_container
    """

    def cmd(self, command, ignore_status=False):
        result = process.run(command, ignore_status=ignore_status, shell=True)
        if result.exit_status == 0:
            return result.exit_status, result.stdout_text.strip()
        else:
            return result.exit_status, result.stderr_text.strip()
        
    def setUp(self):
        super().setUp()
        self.version = self.params.get('registry')
        self.container_engine = self.params.get('engine')
        self.cmd("yum install docker -y --enablerepo=Plus")
        ret_c,containerd_status = self.cmd('systemctl is-active containerd.service',ignore_status=True)
        if containerd_status != 'active':
            self.cmd('systemctl start containerd.service')
        for i in range(30):
            ret_c,docker_status = self.cmd('systemctl is-active docker.service',ignore_status=True)
            if docker_status != 'active':
                time.sleep(20)
                self.cmd('systemctl reset-failed docker.service',ignore_status=True)
                self.cmd('systemctl start docker.service',ignore_status=True)
            else:
                break
        for i in range(5):
            ret_c,pull_log = self.cmd("docker pull %s" % self.version, ignore_status=True)
            if 'Pull complete' in pull_log:
                break
        ret_c,output1 = self.cmd("docker image ls")
        if ':' in self.version:
            self.image_name = self.version.split(':')[0]
        else:
            self.image_name = self.version
        self.assertTrue(self.image_name in output1,"pull image failed")
        ret_c,self.container_id = self.cmd('docker run -dit --privileged=true -p 8888:8888 '+self.version)

    def test(self):
        app_list = ["cuda","libcudnn","libnccl","pytorch","torch-vision","tensorflow2","python3-jupyter-client","python3-jupyter-core","python3-jupyter-events","python3-jupyter-lsp",
                    "python3-jupyter-server","python3-jupyter-server-terminals","python3-jupyterlab","python3-jupyterlab_pygments","python3-jupyterlab-server","noslate-anode"]
        for app in app_list:
            self.cmd("docker exec {} rpm -qa | grep {}".format(self.container_id, app))

    def tearDown(self):
        super().tearDown()
        self.cmd("docker stop %s" % self.container_id)
        self.cmd("docker container prune -f")
        self.cmd("docker image rm %s" % self.version)
        self.cmd("yum remove docker -y")